from enum import Enum


class Model(Enum):
    wanx_style_repaint_v1 = ("人像风格重绘", "wanx-style-repaint-v1")
    image_out_painting = ("图像画面扩展", "image-out-painting")

    def __init__(self, model_cn_name: str, model: str):
        self._model_cn_name = model_cn_name
        self._model = model

    @property
    def model_cn_name(self) -> str:
        return self._model_cn_name

    @property
    def model(self) -> str:
        return self._model

    @classmethod
    def from_model(cls, model: str):
        # 根据模型名称返回对应的枚举值
        for item in cls:
            if item.model == model:
                return item
        raise ValueError(f"No mapping found for model: {model}")


class ModelAPISetting(Enum):
    wanx_style_repaint_v1 = ('aigc', 'image-generation', 'generation')
    image_out_painting = ('aigc', 'image2image', 'out-painting')

    def __init__(self, task_group: str, task: str, function: str):
        self._task_group = task_group
        self._task = task
        self._function = function

    @property
    def task_group(self) -> str:
        return self._task_group

    @property
    def task(self) -> str:
        return self._task

    @property
    def function(self) -> str:
        return self._function

    @classmethod
    def from_model(cls, model: str):
        # 找到与 Model 类对应的枚举值
        model_enum = Model.from_model(model)

        # 这里传递 model_enum 的名字，找到对应的 ModelAPISetting
        for item in cls:
            if item.name == model_enum.name:
                return item
        raise ValueError(f"No API setting found for model: {model}")

